package spaska.data.writers;

import java.io.File;
import java.io.PrintWriter;
import java.util.HashSet;
import java.util.Set;

import spaska.data.Attribute;
import spaska.data.Dataset;
import spaska.data.Instance;
import spaska.data.Value;

public class ARFFWriter implements DataWriter {

	private static final String TAG_RELATION = "@relation";
	private static final String TAG_ATTRIBUT = "@attribute";
	private static final String TAG_DATA = "@data";

	private File file;
	private Set<String> attributes = new HashSet<String>();

	public ARFFWriter(File file, String... attrs) {
		this.file = file;
		for (String a : attrs) {
			attributes.add(a);
		}
	}

	@Override
	public void write(Dataset data) {
		PrintWriter writer = null;

		try {
			writer = new PrintWriter(file);

			writer.println(TAG_RELATION + " " + data.getName());
			writer.println();

			if (attributes.isEmpty()) {
				for (Attribute a : data.getAttributes()) {
					attributes.add(a.getName());
				}
			}
			for (Attribute attr : data.getAttributes()) {
				if (attributes.contains(attr.getName())) {
					String type = null;
					switch (attr.getType()) {
					case Nominal:
						StringBuilder sb = new StringBuilder("{");
						for (Value v : data.getDomain(attr)) {
							sb.append(v.getValue()).append(", ");
						}
						type = sb.substring(0, sb.length() - 2) + "}";
						break;
					case Numeric:
						type = "numeric";
						break;
					default:
						break;
					}
					writer.println(TAG_ATTRIBUT + " " + attr.getName() + " " + type);
				}
			}

			writer.println();
			writer.println(TAG_DATA);
			for (Instance in : data) {
				StringBuilder sb = new StringBuilder();
				int idx = 0;
				for (Attribute at : data.getAttributes()) {
					if (attributes.contains(at.getName())) {
						sb.append(in.valueAt(idx)).append(", ");
					}
					idx++;
				}
				writer.println(sb.substring(0, sb.length() - 2));
			}
		}
		catch (Exception e) {
			e.printStackTrace();
		}
		finally {
			if (writer != null) {
				writer.close();
			}
		}
	}
}
